package com.kingsoft.dc.khaos.module.spark.util

import java.util

import com.alibaba.fastjson.JSONObject
import com.kingsoft.dc.khaos.KhaosContext
import com.kingsoft.dc.khaos.extender.meta.api.{AuthRequest, DmRequest, DmTableSplit, DmTableSplitResult}
import com.kingsoft.dc.khaos.extender.meta.model.MetaParams
import com.kingsoft.dc.khaos.extender.meta.utils.DmMetaUtils
import com.kingsoft.dc.khaos.module.spark.constants.SchedulerConstants
import com.kingsoft.dc.khaos.module.spark.metadata.sink.SubTabelInfo
import com.kingsoft.dc.khaos.module.spark.sink.SinkStrategy
import com.kingsoft.dc.khaos.module.spark.source.SourceStrategy
import org.apache.spark.sql.functions.col
import org.apache.spark.sql.types.StringType
import org.apache.spark.sql.{DataFrame, Dataset, Row}

import scala.collection.JavaConverters._
import scala.collection.mutable

/**
 * Created by haorenhui on 2019/08/10.
 */
object TableSplitUtils {

  //value枚举
  object StrategyTypeEnum extends Serializable {
    val CUSTOM_ENUM = "CUSTOM_ENUM"
    val DATETIME = "DATETIME"
    // 增加业务分表
    val BUSSINESS = "BUSSINESS"
  }

  object StrategyValueEnum extends Serializable {
    val year = "year"
    val month = "month"
    val day = "day"
  }

  def buildQuerySplitTblRequest(kc: KhaosContext): DmRequest = {
    val request = new DmRequest()
    request.setUrl(kc.conf.getString(SchedulerConstants.META_QUERY_URL_KEY))
    val authRequest: AuthRequest = new AuthRequest()
    authRequest.setAuthUrl(kc.conf.getString(SchedulerConstants.AUTH_SERVER_KEY))
    request.setAuthRequest(authRequest)
    request
  }

  def buildCreateSplitTblRequest(kc: KhaosContext): DmRequest = {
    val request = new DmRequest()
    request.setUrl(kc.conf.getString(SchedulerConstants.META_MG_SERVICE_META_MG))
    val authRequest: AuthRequest = new AuthRequest()
    authRequest.setAuthUrl(kc.conf.getString(SchedulerConstants.AUTH_SERVER_KEY))
    request.setAuthRequest(authRequest)
    request
  }

  /** 获取分表信息 */
  def getTableSplit[T](kc: KhaosContext, dbName: String, tblName: String, className: String, metaParamsJson: String): DmTableSplit = {
    val runEnv = kc.conf.getString(SchedulerConstants.RUN_ENV)
    val metaParams = new MetaParams(metaParamsJson)
    val metaRequest = buildQuerySplitTblRequest(kc)
    metaParams.buildTableParams(runEnv, dbName, tblName)
    //metaParams.setProjectName(kc.conf.getString(SchedulerConstants.PROJECT_NAME,null))
    metaRequest.setMetaParams(metaParams)
    val tableSplit: DmTableSplit = DmMetaUtils.getTableSplit(className, metaRequest, classOf[DmTableSplit])
    tableSplit
  }

  /** 获取物理分表名称 */
  def getRealTable[T](kc: KhaosContext, dbName: String, tblName: String, className: String, metaParamsJson: String, classType: T, tableSplit: DmTableSplit, symbol: String, splitValues: scala.List[String]): scala.List[String] = {
    var realTabelList: scala.List[String] = scala.List[String]()

    //获得分表真实名称
    if (classType.isInstanceOf[SourceStrategy]) {
      //读分表
      realTabelList = getReadRealTableList(kc, dbName, tblName, tableSplit, metaParamsJson, symbol, splitValues, className)
    } else if (classType.isInstanceOf[SinkStrategy]) {
      //写分表
      realTabelList = getWriteRealTableList(kc, dbName, tblName, tableSplit, metaParamsJson, symbol, splitValues, className)
    }

    realTabelList
  }

  /** 读取分表时调用 */
  def getReadRealTableList(kc: KhaosContext, dbName: String, tblName: String, tableSplit: DmTableSplit, metaParamsJson: String, symbol: String, splitValues: scala.List[String], className: String): scala.List[String] = {
    var realTabelList: scala.List[String] = scala.List[String]()
    val runEnv = kc.conf.getString(SchedulerConstants.RUN_ENV)
    tableSplit.getStrategyType match {
      //枚举方式分表 TODOO 枚举范围
      case StrategyTypeEnum.CUSTOM_ENUM => {
        //设置dsid dsname dstype projectid db env tbl
        val metaParams = new MetaParams(metaParamsJson)
        metaParams.buildTableParams(runEnv, dbName, tblName)
        //设置请求url ip
        val metaRequest = buildQuerySplitTblRequest(kc)
        metaRequest.setMetaParams(metaParams)
        //设置请求参数的json
        val requestJson = generateQueryTableSplitJson(metaRequest, tableSplit, symbol, splitValues)
        metaRequest.setParams(requestJson)
        realTabelList = DmMetaUtils.getRealTable(className, metaRequest, classOf[DmTableSplitResult]).getRealTblNames.toList
      }
      case StrategyTypeEnum.DATETIME => {
        val metaParams = new MetaParams(metaParamsJson)
        metaParams.buildTableParams(runEnv, dbName, tblName)
        val metaRequest = buildQuerySplitTblRequest(kc)
        metaRequest.setMetaParams(metaParams)
        val requestJson = generateQueryTableSplitJson(metaRequest, tableSplit, symbol, splitValues)
        metaRequest.setParams(requestJson)
        realTabelList = DmMetaUtils.getRealTable(className, metaRequest, classOf[DmTableSplitResult]).getRealTblNames.toList
      }
      // 新增业务分表
      case StrategyTypeEnum.BUSSINESS => {
        //设置dsid dsname dstype projectid db env tbl
        val metaParams = new MetaParams(metaParamsJson)
        metaParams.buildTableParams(runEnv, dbName, tblName)
        //设置请求url ip
        val metaRequest = buildQuerySplitTblRequest(kc)
        metaRequest.setMetaParams(metaParams)
        //设置请求参数的json
        val requestJson = generateQueryTableSplitJson(metaRequest, tableSplit, symbol, splitValues)
        metaRequest.setParams(requestJson)
        realTabelList = DmMetaUtils.getRealTable(className, metaRequest, classOf[DmTableSplitResult]).getRealTblNames.toList
      }
    }
    realTabelList
  }

  /** 写入分表时调用 */
  def getWriteRealTableList(kc: KhaosContext, dbName: String, tblName: String, tableSplit: DmTableSplit, metaParamsJson: String, symbol: String, splitValues: scala.List[String], className: String): scala.List[String] = {
    var realTabelList: scala.List[String] = scala.List[String]()
    val runEnv = kc.conf.getString(SchedulerConstants.RUN_ENV)
    tableSplit.getStrategyType match {
      case StrategyTypeEnum.CUSTOM_ENUM => {
        //设置dsid dsname dstype projectid
        val metaParams = new MetaParams(metaParamsJson)
        metaParams.buildTableParams(runEnv, dbName, tblName)
        //设置请求url ip
        val metaRequest = buildQuerySplitTblRequest(kc)
        metaRequest.setMetaParams(metaParams)
        //设置请求参数的json
        val requestJson = generateQueryTableSplitJson(metaRequest, tableSplit, symbol, splitValues)
        metaRequest.setParams(requestJson)
        realTabelList = DmMetaUtils.getRealTable(className, metaRequest, classOf[DmTableSplitResult]).getRealTblNames.toList
      }
      case StrategyTypeEnum.DATETIME => {
        val metaParams = new MetaParams(metaParamsJson)
        metaParams.buildTableParams(runEnv, dbName, tblName)
        //调用建表的接口
        val metaRequest = buildCreateSplitTblRequest(kc)
        metaRequest.setMetaParams(metaParams)
        val requestJson: String = generateCreateTableSplitJson(metaRequest, splitValues)
        metaRequest.setParams(requestJson)
        realTabelList = DmMetaUtils.getRealTableWithDateTimeAndWrite(className, metaRequest, classOf[DmTableSplitResult]).getRealTblNames.toList
      }
      case StrategyTypeEnum.BUSSINESS => {
        val metaParams = new MetaParams(metaParamsJson)
        metaParams.buildTableParams(runEnv, dbName, tblName)
        //调用建表的接口
        val metaRequest = buildCreateSplitTblRequest(kc)
        metaRequest.setMetaParams(metaParams)
        val requestJson: String = generateCreateTableSplitJson(metaRequest, splitValues)
        metaRequest.setParams(requestJson)
        realTabelList = DmMetaUtils.getRealTableWithDateTimeAndWrite(className, metaRequest, classOf[DmTableSplitResult]).getRealTblNames.toList
      }
    }
    realTabelList
  }

  /** 生成查询分表的请求json */
  def generateQueryTableSplitJson(metaRequest: DmRequest, tableSplit: DmTableSplit, symbol: String, splitValues: scala.List[String]): String = {
    val tblObjects = new util.ArrayList[JSONObject]()
    val searchObjects = new util.ArrayList[JSONObject]()
    val splitRuleObjects = new util.ArrayList[JSONObject]()
    val ruleObject = new JSONObject()
    ruleObject.put("symbol", symbol)
    ruleObject.put("values", splitValues.asJava)
    splitRuleObjects.add(ruleObject)
    val searchObject = new JSONObject()
    searchObject.put("strategyType", tableSplit.getStrategyType)
    searchObject.put("splitColName", tableSplit.getSplitColName)
    searchObject.put("splitRule", splitRuleObjects)
    searchObjects.add(searchObject)
    val jsonObj: JSONObject = new JSONObject()
    jsonObj.put("env", metaRequest.getMetaParams.getEnv)
    jsonObj.put("dsName", metaRequest.getMetaParams.getDsName)
    jsonObj.put("dsId", metaRequest.getMetaParams.getDsId)
    jsonObj.put("dbName", metaRequest.getMetaParams.getDbName)
    jsonObj.put("tlbName", metaRequest.getMetaParams.getTblName)
    jsonObj.put("searchParams", searchObjects)
    val requestJObject = new JSONObject()

    tblObjects.add(jsonObj)
    requestJObject.put("tlbParams", tblObjects)

    requestJObject.toJSONString
  }

  /** 生成创建分表的请求json */
  def generateCreateTableSplitJson(metaRequest: DmRequest, splitValues: scala.List[String]): String = {
    val tblObjects = new util.ArrayList[JSONObject]()
    val requestJObject = new JSONObject()
    val jsonObj: JSONObject = new JSONObject()
    jsonObj.put("env", metaRequest.getMetaParams.getEnv)
    jsonObj.put("dsName", metaRequest.getMetaParams.getDsName)
    jsonObj.put("dbName", metaRequest.getMetaParams.getDbName)
    jsonObj.put("tlbName", metaRequest.getMetaParams.getTblName)
    jsonObj.put("splitValues", splitValues.asJava)
    tblObjects.add(jsonObj)
    requestJObject.put("tlbParams", tblObjects)
    requestJObject.toJSONString
  }


  /** 获得真实物理分表和对应的DF */
  def getSinkRealTable[T](kc: KhaosContext,
                          dbName: String,
                          tblName: String,
                          classType: T,
                          metaClazz: String,
                          metaParamsJson: String,
                          resultDF: DataFrame,
                          subTabelInfo: SubTabelInfo): mutable.HashMap[String, DataFrame] = {
    var returnDF = resultDF
    //获取分表信息
    var tblNameList: scala.List[String] = scala.List[String]()
    val tblNameAndDF: mutable.HashMap[String, DataFrame] = mutable.HashMap[String, DataFrame]()
    val tableSplit: DmTableSplit = TableSplitUtils.getTableSplit(kc, dbName, tblName,
      metaClazz,
      metaParamsJson)
    //不分表
    if (tableSplit == null) {
      //处理不分表的情况下，如果owner表授权过来的，不带库名没有权限。分表模式不需要带库名
      tblNameAndDF.put(tblName, returnDF)
    } else {

      //判断分表方式
      tableSplit.getStrategyType match {
        case StrategyTypeEnum.CUSTOM_ENUM => {
          //根据字段获取分表值
          val splitColName: String = tableSplit.getSplitColName
          val splitValues: List[String] = returnDF.select(col(splitColName).cast(StringType)).distinct().collect().map(row => {
            row.getAs[String](splitColName)
          }).toList.filter(_ != null)
          //获取分表名称列表
          tblNameList = TableSplitUtils.getRealTable(kc, dbName, tblName, metaClazz,
            metaParamsJson, classType, tableSplit, "in", splitValues)
          //生成(分表值,分表名称)map
          val splValues: Array[String] = tableSplit.getStrategyValue.split(",")
          val splitValueAndTblName: Map[String, String] = splValues.map(splitValue => {
            var tp: (String, String) = (splitValue, null)
            tblNameList.foreach(tbl => {
              if (tbl.endsWith("___" + splitValue)) {
                tp = (splitValue, tbl)
              }
            })
            tp
          }).toMap

          //根据分表字段切分DF,并与分表名对应
          returnDF = returnDF.cache()

          for (i <- splValues.indices) {
            val tmpDF: DataFrame = returnDF.filter(col(splitColName) === splValues(i))
            tblNameAndDF.put(splitValueAndTblName(splValues(i)), tmpDF)
          }
          /*val tmpDF: Dataset[Row] = returnDF.filter(!(col(splitColName) isin (splitValues:_*)))
          if (!tmpDF.rdd.isEmpty()) {
              tblNameAndDF.put(tblName + "___others", tmpDF)
          }*/
        }
        case StrategyTypeEnum.DATETIME => {
          //根据字段获取分表值
          val splitColName: String = tableSplit.getSplitColName
          var splitValues: List[String] = returnDF.select(col(splitColName).cast(StringType)).distinct().collect().map(row => {
            row.getAs[String](splitColName)
          }).toList.filter(_ != null)
          //(20190101,2019-01-01) map
          var newSplitValues: Map[String, String] = null
          tableSplit.getStrategyValue match {
            //2019-01-01 ==> 2019/201901/20190101
            case StrategyValueEnum.year => newSplitValues = splitValues.map(sv => (sv.replaceAll("-", "").substring(0, 4), sv)).toMap; splitValues = splitValues.map(_.replaceAll("-", "").substring(0, 4))
            case StrategyValueEnum.month => newSplitValues = splitValues.map(sv => (sv.replaceAll("-", "").substring(0, 6), sv)).toMap; splitValues = splitValues.map(_.replaceAll("-", "").substring(0, 6))
            case StrategyValueEnum.day => newSplitValues = splitValues.map(sv => (sv.replaceAll("-", "").substring(0, 8), sv)).toMap; splitValues = splitValues.map(_.replaceAll("-", "").substring(0, 8))
          }

          //获取分表名称列表
          tblNameList = TableSplitUtils.getRealTable(kc, dbName, tblName, metaClazz,
            metaParamsJson, classType, tableSplit, "in", splitValues)

          //(20190101,table___20190101) map
          val splitValueAndTblName: Map[String, String] = splitValues.map(splitValue => {
            var tp: (String, String) = (splitValue, null)
            tblNameList.foreach(tbl => {
              if (tbl.endsWith("___" + splitValue)) {
                tp = (splitValue, tbl)
              }
            })
            tp
          }).toMap
          returnDF = returnDF.cache()
          //根据分表字段切分DF,并与分表名对应
          for (i <- splitValues.indices) {
            val tmpDF: DataFrame = returnDF.filter(col(splitColName).startsWith(newSplitValues(splitValues(i))))
            tblNameAndDF.put(splitValueAndTblName(splitValues(i)), tmpDF)
          }
          /*//不匹配的数据 放到others表中
          tblNameAndDF.foreach(tp=>returnDF=returnDF.except(tp._2))
          if (!returnDF.rdd.isEmpty()) {
              tblNameAndDF.put(tblName + "___others", returnDF)
          }*/
        }
        case TableSplitUtils.StrategyTypeEnum.BUSSINESS => {
          var suffixValue = "suffix"
          var realTable = tblName
          // 前端获取到的后缀
          if (subTabelInfo.on_off == "true") {
            suffixValue = subTabelInfo.suffix
            //获取分表真实名字
            realTable = tblName + "___" + suffixValue
          } else {
            //获取分表真实名字
            realTable = tblName
          }
          tblNameAndDF.put(realTable, returnDF)
        }
      }
    }
    tblNameAndDF
  }
}
